{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "execution": {},
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NeuromatchAcademy/course-content/blob/master/tutorials/W3D3_OptimalControl/W3D3_Tutorial2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a> &nbsp; <a href=\"https://kaggle.com/kernels/welcome?src=https://raw.githubusercontent.com/NeoNeuron/professional-workshop-3/master/tutorials/W9_OptimalControl/W9_Tutorial2.ipynb\" target=\"_parent\"><img src=\"https://kaggle.com/static/images/open-in-kaggle.svg\" alt=\"Open in Kaggle\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "# Tutorial 2: Optimal Control for Continuous State\n",
    "\n",
    "__Content creators:__ Zhengwei Wu, Shreya Saxena, Xaq Pitkow\n",
    "\n",
    "__Content modified:__ Kai Chen\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "---\n",
    "# Tutorial Objectives\n",
    "In this tutorial, we will implement a continuous control task: you will design control inputs for a linear dynamical system to reach a target state. The state here is continuous-valued, i.e. takes on any real number from $-\\infty$ to $\\infty$.\n",
    "\n",
    "You have already learned about control for binary states in Tutorial 1, and you have learned about stochastic dynamics, latent states, and measurements yesterday. Now we introduce you to the new concepts of designing a controller with full observation of the state (linear qudratic regulator - LQR), and under partial observability of the state (linear quadratic gaussian - LQG).\n",
    "\n",
    "The running example we consider throughout the tutorial is a cat trying to catch a mouse in space, using its handy little jet pack to navigate. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "---\n",
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "execution": {},
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "# Imports\n",
    "\n",
    "import numpy as np\n",
    "import scipy\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import gridspec\n",
    "from math import isclose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@title Figure Settings\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "import ipywidgets as widgets\n",
    "from ipywidgets import interact, fixed, HBox, Layout, VBox, interactive, Label\n",
    "\n",
    "nma_style = {\n",
    "    'figure.figsize' : (8, 6),\n",
    "    'figure.autolayout' : True,\n",
    "    'font.size' : 15,\n",
    "    'xtick.labelsize' : 'small',\n",
    "    'ytick.labelsize' : 'small',\n",
    "    'legend.fontsize' : 'small',\n",
    "    'axes.spines.top' : False,\n",
    "    'axes.spines.right' : False,\n",
    "    'xtick.major.size' : 5,\n",
    "    'ytick.major.size' : 5,\n",
    "}\n",
    "for key, value in nma_style.items():\n",
    "    plt.rcParams[key] = value\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# @title Plotting Functions\n",
    "\n",
    "def plot_vs_time(s, slabel, color, goal=None, ylabel=None):\n",
    "  plt.plot(s, color, label=slabel)\n",
    "  if goal is not None:\n",
    "    plt.plot(goal, 'm', label='goal $g$')\n",
    "  plt.xlabel(\"Time\", fontsize=14)\n",
    "  plt.legend(loc=\"upper right\")\n",
    "\n",
    "  if ylabel:\n",
    "    plt.ylabel(ylabel, fontsize=14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# @title Helper Functions\n",
    "\n",
    "class ExerciseError(AssertionError):\n",
    "  pass\n",
    "\n",
    "\n",
    "def test_lds_class(lds_class):\n",
    "  from math import isclose\n",
    "  ldsys = lds_class(T=2, ini_state=2., noise_var=0.)\n",
    "  if not isclose(ldsys.dynamics(.9)[1], 1.8):\n",
    "      raise ExerciseError(\"'dynamics' method is not correctly implemented!\")\n",
    "  if not isclose(ldsys.dynamics_openloop(.9, 2., np.zeros(ldsys.T)-1.)[1], -0.2):\n",
    "      raise ExerciseError(\"'dynamics_openloop' method is not correctly implemented!\")\n",
    "  if not isclose(ldsys.dynamics_closedloop(.9, 2., np.zeros(ldsys.T)+.3)[0][1], 3.):\n",
    "      raise ExerciseError(\"s[t] in 'dynamics_closedloop' method is not correctly implemented!\")\n",
    "  if not isclose(ldsys.dynamics_closedloop(.9, 2., np.zeros(ldsys.T)+.3)[1][0], .6):\n",
    "      raise ExerciseError(\"a[t] in 'dynamics_closedloop' method is not correctly implemented!\")\n",
    "  ldsys.noise_var = 1.\n",
    "  if isclose(ldsys.dynamics(.9)[1], 1.8):\n",
    "      raise ExerciseError(\"Did you forget to add noise to your s[t+1] in 'dynamics'?\")\n",
    "  if isclose(ldsys.dynamics_openloop(.9, 2., np.zeros(ldsys.T)-1.)[1], -0.2):\n",
    "      raise ExerciseError(\"Did you forget to add noise to your s[t+1] in 'dynamics_openloop'?\")\n",
    "  if isclose(ldsys.dynamics_closedloop(.9, 2., np.zeros(ldsys.T)+.3)[0][1], 3.):\n",
    "      raise ExerciseError(\"Did you forget to add noise to your s[t+1] in 'dynamics_closedloop'?\")\n",
    "  if not isclose(ldsys.dynamics_closedloop(.9, 2., np.zeros(ldsys.T)+.3)[1][0], .6):\n",
    "      raise ExerciseError(\"Your input a[t] should not be noisy in 'dynamics_closedloop'.\")\n",
    "\n",
    "  print('Well Done!')\n",
    "\n",
    "\n",
    "def test_lqr_class(lqr_class):\n",
    "  from math import isclose\n",
    "  lqreg = lqr_class(T=2, ini_state=2., noise_var=0.)\n",
    "  lqreg.goal = np.array([-2, -2])\n",
    "  s = np.array([1, 2])\n",
    "  a = np.array([3, 4])\n",
    "  if not isclose(lqreg.calculate_J_state(s), 25):\n",
    "      raise ExerciseError(\"'calculate_J_state' method is not correctly implemented!\")\n",
    "  if not isclose(lqreg.calculate_J_control(a), 25):\n",
    "      raise ExerciseError(\"'calculate_J_control' method is not correctly implemented!\")\n",
    "\n",
    "  print('Well Done!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "---\n",
    "# Section 1: Exploring a Linear Dynamical System (LDS) with Open-Loop and Closed-Loop Control"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "\n",
    "In this example, a cat is trying to catch a mouse in space. The location of the mouse is the goal state $g$, here a static goal. Later on, we will make the goal time varying, i.e. $g(t)$. The cat's location is the state of the system $s_t$. The state has its internal dynamics: think of the cat drifting slowly in space. These dynamics are such that the state at the next time step $s_{t+1}$ are a linear function of the current state $s_t$. There is some environmental noise (think: meteorites) affecting the state, here modeled as gaussian noise $w_t$.\n",
    "\n",
    "The control input or action $a_t$ is the action of the jet pack, which has an effect $Ba_t$ on the state at the next time step $s_{t+1}$. In this tutorial, we will be designing the action $a_t$ to reach the goal $g$, with known state dynamics.\n",
    "\n",
    "Thus, our linear discrete-time system evolves according to the following equation:\n",
    "\n",
    "\\begin{eqnarray*}\n",
    "s_{t+1} &=& Ds_t + Ba_t + w_t \\tag{1}\\\\\n",
    "s_{0} &=& s_{init}\n",
    "\\end{eqnarray*}\n",
    "\n",
    "with \n",
    "\n",
    "$t$: time step, ranging from $1$ to $T$, where $T$ is the time horizon.\n",
    "\n",
    "$s_t$: state at time $t$ \n",
    "\n",
    "$a_t$: action at time $t$ (also known as control input)\n",
    "\n",
    "$w_t$: gaussian noise at time $t$\n",
    "\n",
    "$D$ and $B$: parameters of the linear dynamical system. \n",
    "\n",
    "For simplicity, we will consider the 1D case, where the matrices reduce to scalars, and the states, control and noise are one-dimensional as well. Specifically, $D$ and $B$ are scalars.\n",
    "\n",
    "We will consider the goal $g$ to be the origin, i.e. $g=0$, for Exercises 1 and 2.2. Note that if the state dynamics are stable, the state reaches $0$ in any case. This is a slightly unrealistic situation for the purposes of simplicity, but we will see more realistic cases later on with $g \\neq 0$\n",
    "\n",
    "**Stability** \\\\\n",
    "The system is stable, i.e. the output remains finite for any finite initial condition $s_{init}$, if $|D|<1$.\n",
    "\n",
    "**Control** \\\\\n",
    "In *open-loop control*, $a_t$ is not a function of $s_t$. In *closed-loop linear control*, $a_t$ is a linear function of the state $s_t$. Specifically, $a_t$ is the control gain $L_t$ multiplied by $s_t$, i.e. $a_t=L_t s_t$. For now, you will explore these equations, and later on, you will *design* $L_t$ to reach the goal $g$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Coding Exercise 1: Implement state evolution equations\n",
    "\n",
    "\n",
    "Implement the state evolution equations in the class methods as provided below, for the following cases: \\\\\n",
    "(a) no control: `def dynamics` \\\\\n",
    "(b) open-loop control: `def dynamics_openloop` \\\\\n",
    "(c) closed-loop control: `def dynamics_closedloop` \\\\\n",
    "\n",
    "*Tip: refer to Equation (1) above. The provided code uses the same notation*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "both",
    "execution": {},
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "class LDS:\n",
    "  def __init__(self, T: int, ini_state: float, noise_var: float):\n",
    "    self.T = T  # time horizon\n",
    "    self.ini_state = ini_state\n",
    "    self.noise_var = noise_var\n",
    "\n",
    "  def dynamics(self, D: float):\n",
    "    s = np.zeros(self.T)  # states initialization\n",
    "    s[0] = self.ini_state\n",
    "    noise = np.random.normal(0, self.noise_var, self.T)\n",
    "\n",
    "    for t in range(self.T - 1):\n",
    "      ####################################################################\n",
    "      ## Insert your code here to fill with the state dynamics equation\n",
    "      ## without any control input\n",
    "      ## complete the function and remove\n",
    "      raise NotImplementedError(\"Exercise: Please complete 'dynamics'\")\n",
    "      ####################################################################\n",
    "      # calculate the state of t+1\n",
    "      s[t + 1] = ...\n",
    "\n",
    "    return s\n",
    "\n",
    "  def dynamics_openloop(self, D: float, B: float, a: np.ndarray):\n",
    "\n",
    "    s = np.zeros(self.T)  # states initialization\n",
    "    s[0] = self.ini_state\n",
    "    noise = np.random.normal(0, self.noise_var, self.T)\n",
    "\n",
    "    for t in range(self.T - 1):\n",
    "      ####################################################################\n",
    "      ## Insert your code here to fill with the state dynamics equation\n",
    "      ## with open-loop control input a[t]\n",
    "      ## complete the function and remove\n",
    "      raise NotImplementedError(\"Please complete 'dynamics_openloop'\")\n",
    "      ####################################################################\n",
    "      # calculate the state of t+1\n",
    "      s[t + 1] = ...\n",
    "\n",
    "    return s\n",
    "\n",
    "  def dynamics_closedloop(self, D: float, B: float, L: np.ndarray):\n",
    "\n",
    "    s = np.zeros(self.T)  # states initialization\n",
    "    s[0] = self.ini_state\n",
    "    noise = np.random.normal(0, self.noise_var, self.T)\n",
    "    a = np.zeros(self.T - 1)\n",
    "\n",
    "    for t in range(self.T - 1):\n",
    "      ####################################################################\n",
    "      ## Insert your code here to fill with the state dynamics equation\n",
    "      ## with closed-loop control input as a function of control gain L.\n",
    "      ## complete the function and remove\n",
    "      raise NotImplementedError(\"Please complete 'dynamics_closedloop'\")\n",
    "      ####################################################################\n",
    "      # calculate the current action\n",
    "      a[t] = ...\n",
    "      # calculate the next state\n",
    "      s[t + 1] = ...\n",
    "\n",
    "    return s, a\n",
    "\n",
    "\n",
    "# Test your function\n",
    "test_lds_class(LDS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "class LDS:\n",
    "  def __init__(self, T: int, ini_state: float, noise_var: float):\n",
    "    self.T = T  # time horizon\n",
    "    self.ini_state = ini_state\n",
    "    self.noise_var = noise_var\n",
    "\n",
    "  def dynamics(self, D: float):\n",
    "    s = np.zeros(self.T)  # states initialization\n",
    "    s[0] = self.ini_state\n",
    "    noise = np.random.normal(0, self.noise_var, self.T)\n",
    "\n",
    "    for t in range(self.T - 1):\n",
    "      # calculate the state of t+1\n",
    "      s[t + 1] = D * s[t] + noise[t]\n",
    "\n",
    "    return s\n",
    "\n",
    "  def dynamics_openloop(self, D: float, B: float, a: np.ndarray):\n",
    "\n",
    "    s = np.zeros(self.T)  # states initialization\n",
    "    s[0] = self.ini_state\n",
    "    noise = np.random.normal(0, self.noise_var, self.T)\n",
    "\n",
    "    for t in range(self.T - 1):\n",
    "      # calculate the state of t+1\n",
    "      s[t + 1] = D * s[t] + B * a[t] + noise[t]\n",
    "\n",
    "    return s\n",
    "\n",
    "  def dynamics_closedloop(self, D: float, B: float, L: np.ndarray):\n",
    "\n",
    "    s = np.zeros(self.T)  # states initialization\n",
    "    s[0] = self.ini_state\n",
    "    noise = np.random.normal(0, self.noise_var, self.T)\n",
    "    a = np.zeros(self.T - 1)\n",
    "\n",
    "    for t in range(self.T - 1):\n",
    "      # calculate the current action\n",
    "      a[t] = L[t] * s[t]\n",
    "      # calculate the next state\n",
    "      s[t + 1] = D * s[t] + B * a[t] + noise[t]\n",
    "\n",
    "    return s, a\n",
    "\n",
    "\n",
    "# Test your function\n",
    "test_lds_class(LDS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Interactive Demo 1.1: Explore no control vs. open-loop control vs. closed-loop control\n",
    "\n",
    "Once your code above passes the tests, use the interactive demo below to visualize the effects of different kinds of control inputs. \n",
    "\n",
    "(a) For the no-control case, can you identify two distinct outcomes, depending on the value of D? Why? \n",
    "\n",
    "(b) The open-loop controller works well--or does it? Run the simulation multiple times and see if there are any problems, especially in challenging (high noise) conditions. \n",
    "\n",
    "(c) Does the closed-loop controller fare better with the noise? Vary the values of $L$ and find a range where it quickly reaches the goal. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "#@markdown Play around (attentively) with **`a`** and **`L`** to see the effect on the open-loop controlled and closed-loop controlled state.\n",
    "\n",
    "def simulate_lds(D=0.95, L=-0.3, a=-1., B=2., noise_var=0.1,\n",
    "                 T=50, ini_state=2.):\n",
    "\n",
    "    # linear dynamical system\n",
    "    lds = LDS(T, ini_state, noise_var)\n",
    "\n",
    "    # No control\n",
    "    s_no_control=lds.dynamics(D)\n",
    "\n",
    "    # Open loop control\n",
    "    at = np.append(a, np.zeros(T - 1))\n",
    "    s_open_loop = lds.dynamics_openloop(D, B, at)\n",
    "\n",
    "    # Closed loop control\n",
    "    Lt =  np.zeros(T) + L\n",
    "    s_closed_loop, a_closed_loop = lds.dynamics_closedloop(D, B, Lt)\n",
    "\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    plt.plot(s_no_control, 'b', label='No control')\n",
    "    plt.plot(s_open_loop, 'g', label='Open Loop with a = {}'.format(a))\n",
    "    plt.plot(s_closed_loop, 'r', label='Closed Loop with L = {}'.format(L))\n",
    "    plt.plot(np.zeros(T), 'm', label='goal')\n",
    "    plt.title('LDS State Evolution')\n",
    "    plt.ylabel('State', fontsize=14)\n",
    "    plt.xlabel('Time', fontsize=14)\n",
    "    plt.legend(loc=\"upper right\")\n",
    "    plt.show()\n",
    "\n",
    "widget=interactive(simulate_lds, {'manual': True},\n",
    "                   D=(.85, 1.05, .1),\n",
    "                   L=(-0.6, 0., .15),\n",
    "                   a=(-2., 1., 1.),\n",
    "                   B=(1., 3., 1.),\n",
    "                   noise_var=(0., 0.2, .1),\n",
    "                   T=fixed(50),\n",
    "                   ini_state=(2., 10., 4.))\n",
    "\n",
    "widget.children[-2].description='Run Simulation'\n",
    "widget.children[-2].style.button_color='lightgreen'\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# to_remove explanation\n",
    "\n",
    "\"\"\"\n",
    "In Exercise 1.2, you should have noticed the following:\n",
    "\n",
    "* No control (blue): the dynamics parameter D controls how fast the dynamics decay\n",
    "  towards 0. For -1<D<1, the system is stable and therefore approaches zero quickly.\n",
    "  However, D>1 produces an unstable system, causing , you should have\n",
    "  noticed that the 'no control' state (blue curve) rapidly explodes\n",
    "  (i.e., heads off to infinity)\n",
    "\n",
    "* Open-loop control: While the open-loop state (green curve) often reachs the goal\n",
    "  quickly, it may not stay there. Under high noise conditions, it tends to\n",
    "  drift away from the goal, though you may not see this in every simulation.\n",
    "\n",
    "* Closed-loop control: The closed-loop state (red curve) reaches the goal and\n",
    "  stays there even in the presence of noise. It converges especially quickly for\n",
    "  Ls around 0.45\n",
    "\n",
    "  Remember that in closed-loop control,\n",
    "  we have a[t]=L[t] * s[t] $. Note that with a constant control gain $L[t]=L,\n",
    "  the state evolution equations can be rearranged to show that the stability of\n",
    "  the closed-loop system now depends on the value of D+BL. (See Equation 2, below).\n",
    "\n",
    "  If $|D+BL|<1$, our closed-loop system will be stable. More generally, you can\n",
    "  view the role of a closed-loop control input as changing the system *dynamics*\n",
    "  in an optimal way to reach the goal.\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Interactive Demo 1.2: Exploring the closed-loop setting further \n",
    "Execute the cell below to visualize the MSE between the state and goal, as a function of control gain $L$. You should see a U-shaped curve, with a clear minimum MSE. The control gain at which the minimum MSE is reached, is the 'optimal' constant control gain for minimizing MSE, here called the numerical optimum. \n",
    "\n",
    "A green dashed line is shown $L = -\\frac{D}{B}$ with $D=0.95$ and $B=2$. Consider how Why is this the theoretical optimal control gain for minimizing MSE of the state $s$ to the goal $g=0$? Examine how the states evolve with a constant gain $L$\n",
    "$$\n",
    "\\begin{eqnarray*}\n",
    "s_{t+1} &=& Ds_t + Ba_t + w_t \\\\\n",
    "&=& Ds_t + B(Ls_t) + w_t \\\\\n",
    "&=& (D+BL)s_t + w_t \\tag{2}\n",
    "\\end{eqnarray*}\n",
    "$$\n",
    "\n",
    "Now, let's visualize the evolution of the system as we change the control gain. We will start with the optimal gain (the control gain that gets us the minimum MSE), and then explore over- and under- ambitious values. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Execute this cell to visualize MSE between state and goal, as a function of control gain\n",
    "def calculate_plot_mse():\n",
    "  D, B, noise_var, T, ini_state = 0.95, 2., 0.1, 50, 2.\n",
    "  control_gain_array = np.linspace(0.1, 0.9, T)\n",
    "  mse_array = np.zeros(control_gain_array.shape)\n",
    "  for i in range(len(control_gain_array)):\n",
    "    lds = LDS(T, ini_state, noise_var)\n",
    "    L = - np.ones(T) * control_gain_array[i]\n",
    "    s, a = lds.dynamics_closedloop(D, B, L)\n",
    "    mse_array[i] = np.sum(s**2)\n",
    "\n",
    "  plt.figure()\n",
    "  plt.plot(-control_gain_array, mse_array, 'b')\n",
    "  plt.axvline(x=-D/B, color='g', linestyle='--')\n",
    "  plt.xlabel(\"control gain (L)\", fontsize=14)\n",
    "  plt.ylabel(\"MSE between state and goal\" , fontsize=14)\n",
    "  plt.title(\"MSE vs control gain\", fontsize=20)\n",
    "  plt.show()\n",
    "\n",
    "\n",
    "calculate_plot_mse()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "#@markdown Explore different values of control gain **`L`** (close to optimal, over- and under- ambitious) \\\\\n",
    "\n",
    "\n",
    "def simulate_L(L:float=-0.45):\n",
    "  D, B, noise_var, T, ini_state = 0.95, 2., 0.1, 50, 2.\n",
    "  lds = LDS(T, ini_state, noise_var)\n",
    "  # Closed loop control with the numerical optimal control gain\n",
    "  Lt = np.ones(T) * L\n",
    "  s_closed_loop_choice, _ = lds.dynamics_closedloop(D, B, Lt)\n",
    "  # Closed loop control with the theoretical optimal control gain\n",
    "  L_theory = - D / B * np.ones(T)\n",
    "  s_closed_loop_theoretical, _ = lds.dynamics_closedloop(D, B, L_theory)\n",
    "  # Plotting closed loop state evolution with both theoretical and numerical optimal control gains\n",
    "  plt.figure(figsize=(10, 6))\n",
    "  plot_vs_time(s_closed_loop_theoretical,\n",
    "                'Closed Loop (Theoretical optimal control gain)','b')\n",
    "  plot_vs_time(s_closed_loop_choice,\n",
    "                'Closed Loop (your choice of L = {})'.format(L), 'g',\n",
    "                goal=np.zeros(T), ylabel=\"State\")\n",
    "  plt.title('Closed Loop State Evolution')\n",
    "  plt.show()\n",
    "\n",
    "\n",
    "widget=interactive(simulate_L, {'manual': True}, L=(-1.05, 0.051, .1))\n",
    "\n",
    "widget.children[-2].description='Run Simulation'\n",
    "widget.children[-2].style.button_color='lightgreen'\n",
    "\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#to_remove explanation\n",
    "\"\"\"\n",
    "In Demo 1.2, you should have seen that the optimal control gain (L = -0.45) takes a\n",
    "short amount of time to get to the goal, and then stays there. We can try to get\n",
    "to the goal in an even shorter time using an 'over-ambitious' control gain (L < -0.45), but\n",
    "this may actually overshoot the goal and may cause oscillations in the system,\n",
    "thus increasing the MSE. On the other hand, an 'under-ambitious' control gain\n",
    "takes a longer time to get to the goal and thus increases the MSE.\n",
    "Finally, at L>0, the system runs away to infinity.\n",
    "\n",
    "Why is L=-D/B optimal for reaching our goal? Recall that our next state is\n",
    "(D+B*L)*s[t] + noise. Plugging that L=-D/B causes that leading term to become zero,\n",
    "which is our goal. Since the noise has zero mean, it's not possible to do any better!\n",
    "\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "---\n",
    "# Section 2: Designing an optimal control input using a linear quadratic regulator (LQR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "## Section 2.1 Constraints on the system\n",
    "Now we will start imposing additional constraints on our system. For example. \n",
    "if you explored different values for $s_{init}$ above, you would have seen very large values for $a_t$ in order to get to the mouse in a short amount of time. However, perhaps the design of our jetpack makes it dangerous to use large amounts of fuel in a single timestep. We certainly do not want to explode, so we would like to keep the actions $a_t$ as small as possible while still mantaining good control.\n",
    "\n",
    "Moreover, in Exercise 1, we had restricted ourselves to a static control gain $L_t \\equiv L$. How would we vary it if we could?\n",
    "\n",
    "This leads us to a more principled way of designing the optimal control input.\n",
    "### Setting up a cost function \n",
    "\n",
    "In a finite-horizon LQR problem,  the cost function is defined as: \n",
    "\n",
    "\\begin{eqnarray}\n",
    "J({\\bf s},{\\bf a}) &=& J_{state}({\\bf s}) + \\rho J_{control}({\\bf a}) \\\\\n",
    " &=& \\sum_{t = 0}^{T} (s_{t}-g)^2 + \\rho \\sum_{t=0}^{T-1}a_{t}^2 \\tag{3}\n",
    "\\end{eqnarray}\n",
    "\n",
    "where $\\rho$ is the weight on the control effort cost, as compared to the cost of not being at the goal. Here, ${\\bf a} = \\{a_t\\}_{t=0}^{T-1}$, ${\\bf s} = \\{s_t\\}_{t=0}^{T}$. This is a quadratic cost function. In Exercise $2$, we will only explore $g=0$, in which case $J_{state}({\\bf s})$ can also be expressed as $\\sum_{t = 0}^{T} s_{t}^2$. In Exercise $3$, we will explore a non-zero time-varying goal.\n",
    "\n",
    "The goal of the LQR problem is to find control ${\\bf a}$ such that $J({\\bf s},{\\bf a})$ is minimized. The goal is then to find the control gain at each time point, i.e.,\n",
    "\n",
    "$$ \\text{argmin} _{\\{L_t\\}_{t=0}^{T-1}}  J({\\bf s},{\\bf a}) \\tag{4} $$ \n",
    "\n",
    "where $a_t = L_t s_t$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "\n",
    "## Section 2.2 Solving LQR\n",
    "The solution to Equation (4), i.e. LQR for a finite time horizon, can be obtained via Dynamic Programming. For details, check out [this lecture by Stephen Boyd](https://stanford.edu/class/ee363/lectures/dlqr.pdf).\n",
    "\n",
    "For an infinite time horizon, one can obtain a closed-form solution using Riccati equations, and the solution for the control gain becomes time-invariant, i.e. $L_t \\equiv L$. We will use this in Exercise 4. For  details, check out [this other lecture by Stephen Boyd](https://stanford.edu/class/ee363/lectures/dlqr-ss.pdf).\n",
    "\n",
    "Additional reference for entire section: \\\\\n",
    "[Bertsekas, Dimitri P. \"Dynamic programming and optimal control\". Vol. 1. No. 2. Belmont, MA: Athena scientific, 1995.](http://www.athenasc.com/dpbook.html)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Coding Exercise 2.2: Implement the cost function\n",
    "The cost function $J_{control}({\\bf s}, {\\bf a})$ can be divided into two parts: $J_{state}({\\bf s})$ and $J_{control}({\\bf a})$. \n",
    "\n",
    "Code up these two parts in the class methods `def calculate_J_state` and `def calculate_J_control` in the following helper class for LQR. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "class LQR(LDS):\n",
    "  def __init__(self, T, ini_state, noise_var):\n",
    "    super().__init__(T, ini_state, noise_var)\n",
    "    self.goal = np.zeros(T)  # The class LQR only supports g=0\n",
    "\n",
    "  def control_gain_LQR(self, D, B, rho):\n",
    "    P = np.zeros(self.T)  # Dynamic programming variable\n",
    "    L = np.zeros(self.T - 1)  # control gain\n",
    "    P[-1] = 1\n",
    "\n",
    "    for t in range(self.T - 1):\n",
    "        P_t_1 = P[self.T - t - 1]\n",
    "        P[self.T - t-2] = (1 + P_t_1 * D**2 - D * P_t_1 * B / (\n",
    "                rho + P_t_1 * B) * B**2 * P_t_1 * D)\n",
    "\n",
    "        L[self.T - t-2] = - (1 / (rho + P_t_1 * B**2) * B * P_t_1 * D)\n",
    "    return L\n",
    "\n",
    "  def calculate_J_state(self, s:np.ndarray):\n",
    "    ########################################################################\n",
    "    ## Insert your code here to calculate J_state(s) (see Eq. 3)\n",
    "    ## complete the function and remove\n",
    "    raise NotImplementedError(\"Please complete 'calculate_J_state'\")\n",
    "    ########################################################################\n",
    "    # calculate the state\n",
    "    J_state = ...\n",
    "\n",
    "    return J_state\n",
    "\n",
    "  def calculate_J_control(self, a:np.ndarray):\n",
    "    ########################################################################\n",
    "    ## Insert your code here to calculate J_control(a) (see Eq. 3).\n",
    "    ## complete the function and remove\n",
    "    raise NotImplementedError(\"Please complete 'calculate_J_control'\")\n",
    "    ########################################################################\n",
    "    # calculate the control\n",
    "    J_control = ...\n",
    "\n",
    "    return J_control\n",
    "\n",
    "\n",
    "# Test class\n",
    "test_lqr_class(LQR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# to_remove solution\n",
    "class LQR(LDS):\n",
    "  def __init__(self, T, ini_state, noise_var):\n",
    "    super().__init__(T, ini_state, noise_var)\n",
    "    self.goal = np.zeros(T)  # The class LQR only supports g=0\n",
    "\n",
    "  def control_gain_LQR(self, D, B, rho):\n",
    "    P = np.zeros(self.T)  # Dynamic programming variable\n",
    "    L = np.zeros(self.T - 1)  # control gain\n",
    "    P[-1] = 1\n",
    "\n",
    "    for t in range(self.T - 1):\n",
    "        P_t_1 = P[self.T - t - 1]\n",
    "        P[self.T - t-2] = (1 + P_t_1 * D**2 - D * P_t_1 * B / (\n",
    "                rho + P_t_1 * B) * B**2 * P_t_1 * D)\n",
    "\n",
    "        L[self.T - t-2] = - (1 / (rho + P_t_1 * B**2) * B * P_t_1 * D)\n",
    "    return L\n",
    "\n",
    "  def calculate_J_state(self, s:np.ndarray):\n",
    "    # calculate the state\n",
    "    J_state = np.sum((s - self.goal)**2)\n",
    "\n",
    "    return J_state\n",
    "\n",
    "  def calculate_J_control(self, a:np.ndarray):\n",
    "    # calculate the control\n",
    "    J_control = np.sum(a**2)\n",
    "\n",
    "    return J_control\n",
    "\n",
    "# Test class\n",
    "test_lqr_class(LQR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Interactive Demo 2.2: LQR to the origin \n",
    "\n",
    "In this exercise, we will use your new LQR controller to track a static goal at $g=0$. Here, we will explore how varying $\\rho$ affects its actions by\\\\\n",
    "\n",
    "1.   Using Equation 3, find a value for $\\rho$ that will get you the same cost and control gain as Exercise 1.\n",
    "2.   Pick a larger value for $\\rho$ and see the effect on the action.\n",
    "3.   Try increasing the rho to 2. What do you notice? \\\\\n",
    "4.   For different values of $\\rho$, how does the control gain vary? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "def simulate_rho(rho=1.):\n",
    "    D, B, T, ini_state, noise_var = 0.9, 2., 50, 2., .1  # state parameter\n",
    "    lqr = LQR(T, ini_state, noise_var)\n",
    "    L = lqr.control_gain_LQR(D, B, rho)\n",
    "    s_lqr, a_lqr = lqr.dynamics_closedloop(D, B, L)\n",
    "\n",
    "    plt.figure(figsize=(14, 4))\n",
    "    plt.suptitle('LQR Control for rho = {}'.format(rho), y=1.05)\n",
    "\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plot_vs_time(s_lqr,'State evolution','b',goal=np.zeros(T))\n",
    "    plt.ylabel('State $s_t$')\n",
    "\n",
    "    plt.subplot(1, 3, 2)\n",
    "    plot_vs_time(a_lqr,'LQR Action','b')\n",
    "    plt.ylabel('Action $a_t$')\n",
    "\n",
    "    plt.subplot(1, 3, 3)\n",
    "    plot_vs_time(L,'Control Gain','b')\n",
    "    plt.ylabel('Control Gain $L_t$')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "widget=interactive(simulate_rho, {'manual': True}, rho=(0., 2., 0.5))\n",
    "\n",
    "widget.children[-2].description = 'Run Simulation'\n",
    "widget.children[-2].style.button_color = 'lightgreen'\n",
    "\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# to_remove explanation\n",
    "\"\"\"\n",
    "* rho=0 will get you the same cost and control gain as in Exercise 1 by\n",
    "  zeroing out the cost term.\n",
    "* A small value for rho will have a similar solution as in (a), but with\n",
    "   potentially large values for |a[t]|.\n",
    "* A large value for rho, like 2 will lead to small values for |a[t]|.\n",
    "* The control gain becomes more time-varying (as opposed to fairly static)\n",
    "  for large rho. For some parameter values, L[t] oscillates during the entire\n",
    "  trajectory in order to keep $|a_t|$ low. Try D = 0.9, B = 2 and rho = 2.\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "## Section 2.3: The tradeoff between state cost and control cost\n",
    "\n",
    "In Exercise 2.1, you implemented code to calculate for $J_{state}$ and $J_{control}$ in the class methods for the class LQR. \n",
    "\n",
    "We will now plot them against each other for varying values of $\\rho$ to explore the tradeoff between state cost and control cost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Execute this cell to visualize the tradeoff between state and control cost\n",
    "def calculate_plot_costs():\n",
    "  D, B, noise_var, T, ini_state = 0.9, 2., 0.1, 50, 2.\n",
    "  rho_array = np.linspace(0.2, 40, 100)\n",
    "  J_state = np.zeros(rho_array.shape)\n",
    "  J_control = np.zeros(rho_array.shape)\n",
    "  for i in np.arange(len(rho_array)):\n",
    "    lqr = LQR(T, ini_state, noise_var)\n",
    "    L = lqr.control_gain_LQR(D, B, rho_array[i])\n",
    "    s_lqr, a_lqr = lqr.dynamics_closedloop(D, B, L)\n",
    "    J_state[i] = lqr.calculate_J_state(s_lqr)\n",
    "    J_control[i] = lqr.calculate_J_control(a_lqr)\n",
    "\n",
    "  fig = plt.figure(figsize=(6, 6))\n",
    "  plt.plot(J_state, J_control, '.b')\n",
    "  plt.xlabel(\"$J_{state} = \\sum_{t = 0}^{T} (s_{t}-g)^2$\", fontsize=14)\n",
    "  plt.ylabel(\"$J_{control} = \\sum_{t=0}^{T-1}a_{t}^2$\" , fontsize=14)\n",
    "  plt.title(\"Error vs control effort\", fontsize=20)\n",
    "  plt.show()\n",
    "\n",
    "\n",
    "calculate_plot_costs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "You should notice the bottom half of a 'C' shaped curve, forming the tradeoff between the state cost and the control cost under optimal linear control.\n",
    "For a desired value of the state cost, we cannot reach a lower control cost than the curve in the above plot. Similarly, for a desired value of the control cost, we must accept that amount of state cost. For example, if you know that you have a limited amount of fuel, which determines your maximum control cost to be $J_{control}^{max}$. \n",
    "\n",
    "You will be able to show that you will not be able to track your state with a higher accuracy than the corresponding $J_{state}$ as given by the graph above. This is thus an important curve when designing a system and exploring its control.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "---\n",
    "# Section 3: LQR for tracking a time-varying goal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "In a more realistic situation, the mouse would move around constantly. Suppose you were able to predict the movement of the mouse as it bounces from one place to another. This becomes your goal trajectory $g_t$.\n",
    "\n",
    "When the target state, denoted as $g_t$, is not $0$, the cost function becomes\n",
    "$$ J({\\bf a}) = \\sum_{t = 0}^{T} (s_{t}- g_t) ^2 + \\rho \\sum_{t=0}^{T-1}(a_{t}-\\bar a_t)^2$$\n",
    "Here, $\\bar a_t$ is the desired action based on the goal trajectory. In other words, the controller considers the goal for the next time step, and designs a preliminary control action that gets the state at the next time step to the desired goal. Specifically, without taking into account noise $w_t$, we would like to design $\\bar a_t$ such that $s_{t+1}=g_{t+1}$. Thus, from Equation $(1)$,\n",
    "\n",
    "\\begin{eqnarray*}\n",
    "g_{t+1} &=& Ds_t + B \\bar a_t\\\\\n",
    "\\bar a_{t} &=& \\frac{- Ds_t + g_{t+1}}{B}\\\\\n",
    "\\end{eqnarray*}\n",
    "\n",
    "The final control action $a_t$ is produced by adding this desired action $\\bar a_t$ with the term with the control gain $L_t(s_t - g_t)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Execute this cell to include class\n",
    "#@markdown for LQR control to desired time-varying goal\n",
    "\n",
    "class LQR_tracking(LQR):\n",
    "  def __init__(self, T, ini_state, noise_var, goal):\n",
    "    super().__init__(T, ini_state, noise_var)\n",
    "    self.goal = goal\n",
    "\n",
    "  def dynamics_tracking(self, D, B, L):\n",
    "\n",
    "    s = np.zeros(self.T) # states intialization\n",
    "    s[0] = self.ini_state\n",
    "\n",
    "    noise = np.random.normal(0, self.noise_var, self.T)\n",
    "\n",
    "    a = np.zeros(self.T) # control intialization\n",
    "    a_bar = np.zeros(self.T)\n",
    "    for t in range(self.T - 1):\n",
    "        a_bar[t] = ( - D * s[t] + self.goal[t + 1]) / B\n",
    "        a[t] =  L[t] * (s[t] - self.goal[t]) + a_bar[t]\n",
    "        s[t + 1] = D * s[t] + B * a[t] + noise[t]\n",
    "\n",
    "    return s, a, a_bar\n",
    "\n",
    "  def calculate_J_state(self,s):\n",
    "    J_state = np.sum((s-self.g)**2)\n",
    "    return J_state\n",
    "\n",
    "  def calculate_J_control(self, a, a_bar):\n",
    "    J_control = np.sum((a-a_bar)**2)\n",
    "    return J_control"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Interactive Demo 3: LQR control to desired time-varying goal\n",
    "Use the demo below to explore how LQR tracks a time-varying goal. \n",
    "Starting with the sinusoidal goal function `sin`, investigate how the system reacts with different values of $\\rho$ and process noise variance. Next, explore other time-varying goal, such as a step function and ramp."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "def simulate_tracking(rho=20., noise_var=0.1, goal_func='sin'):\n",
    "  D, B, T, ini_state = 0.9, 2., 100, 0.\n",
    "  if goal_func == 'sin':\n",
    "      goal = np.sin(np.arange(T) * 2 * np.pi * 5 / T)\n",
    "  elif goal_func == 'step':\n",
    "      goal = np.zeros(T)\n",
    "      goal[int(T / 3):] = 1.\n",
    "  elif goal_func == 'ramp':\n",
    "      goal = np.zeros(T)\n",
    "      goal[int(T / 3):] = np.arange(T - int(T / 3)) / (T - int(T / 3))\n",
    "\n",
    "  lqr_time = LQR_tracking(T, ini_state, noise_var, goal)\n",
    "  L = lqr_time.control_gain_LQR(D, B, rho)\n",
    "  s_lqr_time, a_lqr_time, a_bar_lqr_time = lqr_time.dynamics_tracking(D, B, L)\n",
    "\n",
    "  plt.figure(figsize=(13, 5))\n",
    "  plt.suptitle('LQR Control for time-varying goal', y=1.05)\n",
    "  plt.subplot(1, 2, 1)\n",
    "  plot_vs_time(s_lqr_time,'State evolution $s_t$','b',goal, ylabel=\"State\")\n",
    "  plt.subplot(1, 2, 2)\n",
    "  plot_vs_time(a_lqr_time, 'Action $a_t$', 'b', ylabel=\"Action\")\n",
    "  plt.show()\n",
    "\n",
    "\n",
    "widget=interactive(simulate_tracking, {'manual': True},\n",
    "                   rho=(0., 40., 10.),\n",
    "                   noise_var=(0., 1., .2),\n",
    "                   goal_func=['sin', 'step', 'ramp']\n",
    "                   )\n",
    "widget.children[-2].description = 'Run Simulation'\n",
    "widget.children[-2].style.button_color = 'lightgreen'\n",
    "\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# to_remove explanation\n",
    "\n",
    "\"\"\"\n",
    "In Exercise 3, you should have noticed that:\n",
    "* The system follows time varying goals rather well, with little change to the\n",
    "   cost function and the control equations.\n",
    "\n",
    "* Setting rho=0 leads to noise in the first part of the time series.\n",
    "  Here, we see that the control cost in fact acts as a regularizer.\n",
    "\n",
    "* Larger values of the process noise variance lead to a higher MSE between the\n",
    "  state and the desired goal.\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "---\n",
    "# Section 4: Control of an partially observed state using a Linear Quadratic Gaussian (LQG) controller\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "## Section 4.1 Introducing the LQG Controller"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "In practice, the controller does not have full access to the state. For example, your jet pack in space may be controlled by Mission Control back on earth!  In this case, noisy measurements $m_t$ of the state $s_t$ are taken via radar, and the controller needs to (1) estimate the true state, and (2) design an action based on this estimate. \n",
    "\n",
    "Fortunately, the separation principle tells us that it is optimal to do (1) and (2) separately. This makes our problem much easier, since we already know how to do each step.  \n",
    "\n",
    "1) *State Estimation*  \n",
    "Can we recover the state from the measurement? \n",
    "From yesterday's lecture, it is known that the states $\\hat{s}_t$ can be estimated from the measurements $m_t$ using the __Kalman filter__. \n",
    "\n",
    "2) *Design Action*  \n",
    "In Sections 2 and 3 above, we just learnt about the LQR controller which designs an action based on the state. The separation principle tells us that it is sufficient to replace the use of the state in LQR with the *estimated* state, i.e.\n",
    "\n",
    "$$a_t = L_t \\hat s_t$$\n",
    "\n",
    "The state dynamics will then be:\n",
    "$$s_{t+1} = D s_t + B a_t + w_t$$\n",
    "where $w_t$ is the process noise (proc_noise), and the observation / measurement is:\n",
    "$$ y_t = C s_t + v_t$$ \n",
    "with $v_t$ being the measurement noise (meas_noise).\n",
    "\n",
    "The combination of (1) state estimation and (2) action design using LQR is known as a **linear quadratic gaussian (LQG)**. Yesterday, you completed the code for Kalman filter.  Based on that, you will code up the LQG controller. For these exercises, we will resturn to using the goal $g=0$, as in Section 2.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Interactive Demo 4.1: The Kalman filter in conjunction with a linear closed-loop controller (LQG Control)\n",
    "In the `MyKalmanFilter` class, the method `filter_control` implements filtering in closed-loop feedback. It is a combination of generating samples (states $s_t$) and filtering (generating state estimates $\\hat s_t$), as you have seen in yesterday's tutorial. The only difference from yesterday is that today's Kalman filter is in closed loop with the controller. Thus, each $s_{t+1}$ gets an input $a_t$, which itself depends on the state estimate of the last time step $\\hat s_t$.\n",
    "\n",
    "Below you find the code snipets for the Kalman filter in closed loop (`MyKalmanFilter`) class that provide you an insight in action update (`control_policy_LQG`) and state estimation (`state_dynamics_LQG`). Please feel free to inspect the helper functions and classes for the details.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "You should have seen the next cell containing `MyKalmanFilter` class yesterday, with the exception of the controller acting on the state estimate in feedback, using the methods/equations you will find below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Execute this cell to include MyKalmanFilter class\n",
    "\n",
    "class MyKalmanFilter():\n",
    "  def __init__(self, n_dim_state, n_dim_obs, transition_matrices, transition_covariance, observation_matrices,\n",
    "              observation_covariance, initial_state_mean, initial_state_covariance, control_matrices):\n",
    "    \"\"\"\n",
    "    @param n_dim_state: dimension of the latent variables\n",
    "    @param n_dim_obs: dimension of the observed variables\n",
    "    @param transition_matrices: D\n",
    "    @param transition_covariance: process noise\n",
    "    @param observation_matrices: C\n",
    "    @param observation_covariance: measurement noise\n",
    "    @param initial_state_mean: initial state estimate\n",
    "    @param initial_state_covariance: initial estimate on state variance\n",
    "    @param control_matrices: B\n",
    "    \"\"\"\n",
    "    self.n_dim_state = n_dim_state\n",
    "    self.n_dim_obs = n_dim_obs\n",
    "    self.transition_matrices = transition_matrices\n",
    "    self.transition_covariance = transition_covariance\n",
    "    self.observation_matrices = observation_matrices\n",
    "    self.observation_covariance = observation_covariance\n",
    "    self.initial_state_mean = initial_state_mean\n",
    "    self.initial_state_covariance = initial_state_covariance\n",
    "\n",
    "    self.control_matrices = control_matrices\n",
    "\n",
    "  def filter_control(self, n_timesteps, control_gain, use_myfilter=True):\n",
    "    \"\"\"\n",
    "    Method that performs Kalman filtering with a controller in feedback\n",
    "    @param n_timesteps: length of the data sample\n",
    "    @param control_gain: a numpy array whose dimension is [n_timesteps, self.n_dim_state]\n",
    "    @output: filtered_state_means: a numpy array whose dimension is [n_timesteps, self.n_dim_state]\n",
    "    @output: filtered_state_covariances: a numpy array whose dimension is [n_timesteps, self.n_dim_state, self.n_dim_state]\n",
    "    @output: latent_state: a numpy array whose dimension is [n_timesteps, self.n_dim_state]\n",
    "    @output: observed_state: a numpy array whose dimension is [n_timesteps, self.n_dim_obs]\n",
    "    @output: control: a numpy array whose dimension is [n_timesteps, self.n_dim_state]\n",
    "    \"\"\"\n",
    "\n",
    "    # validate inputs\n",
    "    # assert observed_dim == self.n_dim_obs\n",
    "\n",
    "    n_example = n_timesteps\n",
    "    observed_dim = self.n_dim_obs\n",
    "    latent_state = []\n",
    "    observed_state = []\n",
    "    control = []\n",
    "\n",
    "    current_latent_state = self.initial_state_mean #initial_state\n",
    "    control.append(self.initial_state_mean)\n",
    "    latent_state.append(current_latent_state)\n",
    "    observed_state.append(np.dot(self.observation_matrices, current_latent_state) +\n",
    "                              np.random.multivariate_normal(np.zeros(self.n_dim_obs), self.observation_covariance))\n",
    "\n",
    "\n",
    "\n",
    "    # create holders for outputs\n",
    "    filtered_state_means = np.zeros([n_example, self.n_dim_state])\n",
    "    filtered_state_covariances = np.zeros([n_example, self.n_dim_state, self.n_dim_state])\n",
    "\n",
    "\n",
    "    if use_myfilter:\n",
    "      # the first state mean and state covar is the initial expectation\n",
    "      filtered_state_means[0] = self.initial_state_mean\n",
    "      filtered_state_covariances[0] = self.initial_state_covariance\n",
    "\n",
    "      # initialize internal variables\n",
    "      current_state_mean = self.initial_state_mean.copy()\n",
    "      current_state_covar = self.initial_state_covariance.copy()\n",
    "      self.p_n_list = np.zeros((n_example, self.n_dim_obs, self.n_dim_obs))\n",
    "\n",
    "      for i in range(1, n_example):\n",
    "        ## Use the code in Exercise 4.1 to get the current action\n",
    "        current_action = control_policy_LQG(self,current_state_mean,control_gain[i])\n",
    "        control.append(current_action)\n",
    "\n",
    "\n",
    "        ## Use the code in Exercise 4.1 to update the state\n",
    "        current_latent_state = state_dynamics_LQG(self,current_latent_state, current_action)\n",
    "        latent_state.append(current_latent_state)\n",
    "\n",
    "        # use observation_matrices and observation_covariance to calculate next observed state\n",
    "        observed_state.append(np.dot(self.observation_matrices, current_latent_state\n",
    "                                    ) + np.random.multivariate_normal(np.zeros(self.n_dim_obs), self.observation_covariance))\n",
    "\n",
    "        current_observed_data = observed_state[-1]\n",
    "\n",
    "\n",
    "        # run a single step forward filter\n",
    "        # prediction step\n",
    "\n",
    "        predicted_state_mean = np.dot(self.transition_matrices, current_state_mean\n",
    "                                      ) + np.dot(self.control_matrices, current_action)\n",
    "        predicted_state_cov = np.matmul(np.matmul(self.transition_matrices, current_state_covar),\n",
    "                                        np.transpose(self.transition_matrices)) + self.transition_covariance\n",
    "        # observation step\n",
    "        innovation = current_observed_data - np.dot(self.observation_matrices, predicted_state_mean)\n",
    "        innovation_covariance = np.matmul(np.matmul(self.observation_matrices, predicted_state_cov),\n",
    "                                          np.transpose(self.observation_matrices)) + self.observation_covariance\n",
    "        # update step\n",
    "\n",
    "\n",
    "        kalman_gain = np.matmul(np.matmul(predicted_state_cov, np.transpose(self.observation_matrices)),\n",
    "                                np.linalg.inv(innovation_covariance))\n",
    "        current_state_mean = predicted_state_mean + np.dot(kalman_gain, innovation)\n",
    "        current_state_covar = np.matmul((np.eye(current_state_covar.shape[0]) -\n",
    "                                          np.matmul(kalman_gain, self.observation_matrices)),\n",
    "                                        predicted_state_cov)\n",
    "        # populate holders\n",
    "        filtered_state_means[i, :] = current_state_mean\n",
    "        filtered_state_covariances[i, :, :] = current_state_covar\n",
    "        self.p_n_list[i, :, :] = predicted_state_cov\n",
    "          # self.p_n_list[i-1, :, :] = predicted_state_cov\n",
    "          # new\n",
    "          # self.p_n_list[-1, :, :] = np.matmul(np.matmul(self.transition_matrices, filtered_state_covariances[-1,:,:]),\n",
    "          #                                    np.linalg.inv(self.transition_matrices)) + self.transition_covariance\n",
    "\n",
    "#         else:\n",
    "#             #################################################################################\n",
    "#             # below: this is an alternative if you do not have an implementation of filtering\n",
    "#             kf = KalmanFilter(n_dim_state=self.n_dim_state, n_dim_obs=self.n_dim_obs)\n",
    "#             need_params = ['transition_matrices', 'observation_matrices', 'transition_covariance',\n",
    "#                            'observation_covariance', 'initial_state_mean', 'initial_state_covariance']\n",
    "#             for param in need_params:\n",
    "#                 setattr(kf, param, getattr(self, param))\n",
    "#             filtered_state_means, filtered_state_covariances = kf.filter(X)\n",
    "#             #################################################################################\n",
    "\n",
    "    filtered_state_means = np.squeeze(np.array(filtered_state_means))\n",
    "    filtered_state_covariances = np.squeeze(np.array(filtered_state_covariances))\n",
    "    latent_state = np.squeeze(np.array(latent_state))\n",
    "    observed_state = np.squeeze(np.array(observed_state))\n",
    "    control = np.squeeze(np.array(control))\n",
    "\n",
    "\n",
    "    return filtered_state_means, filtered_state_covariances, latent_state, observed_state, control\n",
    "\n",
    "  def plot_state_vs_time(self, n_timesteps, control_gain, title, use_myfilter=True, goal=None):\n",
    "    filtered_state_means_impl, filtered_state_covariances_impl, latent, measurement, control = self.filter_control(\n",
    "        n_timesteps, control_gain)\n",
    "\n",
    "    fig = plt.figure(figsize=(12, 4))\n",
    "    plt.suptitle(title, y=1.05)\n",
    "    gs = gridspec.GridSpec(1, 2, width_ratios=[1, 2])\n",
    "\n",
    "    ax0 = plt.subplot(gs[0])\n",
    "    ax0.plot(latent,filtered_state_means_impl, 'b.')\n",
    "    ax0.set_xlabel('Latent State')\n",
    "    ax0.set_ylabel('Estimated State')\n",
    "    ax0.set_aspect('equal')\n",
    "\n",
    "    ax1 = plt.subplot(gs[1])\n",
    "    ax1.plot(latent, 'b', label = 'Latent State')\n",
    "    ax1.plot(filtered_state_means_impl, 'r', label = 'Estimated State')\n",
    "    if goal is not None:\n",
    "      ax1.plot(goal, 'm', label = 'goal')\n",
    "    ax1.set_xlabel('Time')\n",
    "    ax1.set_ylabel('State')\n",
    "    ax1.legend(loc=\"upper right\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "# inspect the 'control_policy_LQG' and 'state_dynamics_LQG' methods:\n",
    "\n",
    "def control_policy_LQG(self, mean_estimated_state, control_gain):\n",
    "  current_action =  control_gain * mean_estimated_state\n",
    "  return current_action\n",
    "\n",
    "def state_dynamics_LQG(self, current_latent_state, current_action):\n",
    "\n",
    "  current_latent_state = np.dot(self.transition_matrices, current_latent_state)\\\n",
    "                        + np.dot(self.control_matrices, current_action)\\\n",
    "                        + np.random.multivariate_normal(np.zeros(self.n_dim_state),\n",
    "                                                        self.transition_covariance)\n",
    "  return current_latent_state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "Take a look at the helper code for the `MyKalmanFilter` class above. In the following exercises, we will use the same notation that we have been using in this tutorial; adapter code has been provided to convert it into the representation `MyKalmanFilter expects`.\n",
    "\n",
    "Use interactive demo below to refresh your memory of how a Kalman filter estimates state. `C` scales the observation matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "def simulate_kf_no_control(D=0.9, B=2., C=1., L=0., T=50, ini_state=5,\n",
    "                           proc_noise = 0.1, meas_noise = 0.2):\n",
    "\n",
    "  control_gain = np.ones(T) * L\n",
    "\n",
    "  # Format the above variables into a format acccepted by the Kalman Filter\n",
    "  n_dim_state = 1\n",
    "  n_dim_obs = 1\n",
    "  n_timesteps = T\n",
    "\n",
    "  transition_matrices = np.eye(n_dim_state) * D\n",
    "  transition_covariance = np.eye(n_dim_obs) * proc_noise # process noise\n",
    "  observation_matrices = np.eye(n_dim_state) * C\n",
    "  observation_covariance = np.eye(n_dim_obs) * meas_noise\n",
    "  initial_state_mean = np.ones(n_dim_state) * ini_state\n",
    "  initial_state_covariance = np.eye(n_dim_state) * .01\n",
    "  control_matrices = np.eye(n_dim_state) * B\n",
    "\n",
    "  my_kf = MyKalmanFilter(n_dim_state, n_dim_obs, transition_matrices,\n",
    "                      transition_covariance, observation_matrices,\n",
    "                      observation_covariance, initial_state_mean,\n",
    "                      initial_state_covariance, control_matrices)\n",
    "\n",
    "  my_kf.plot_state_vs_time(n_timesteps, control_gain,\n",
    "                            'State estimation with KF (no control input)')\n",
    "\n",
    "\n",
    "widget=interactive(simulate_kf_no_control, {'manual': True},\n",
    "                   D=fixed(.95),\n",
    "                   B=fixed(2.),\n",
    "                   C=(0., 3., 1.),\n",
    "                   proc_noise=(0., 1., .1),\n",
    "                   meas_noise=(0.1, 1., .1),\n",
    "                   T=fixed(50),\n",
    "                   L=fixed(0),\n",
    "                   ini_state=fixed(5.))\n",
    "\n",
    "widget.children[-2].description = 'Run Simulation'\n",
    "widget.children[-2].style.button_color = 'lightgreen'\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#to_remove explanation\n",
    "\"\"\"\n",
    "You should have seen that the Kalman filter generally estimates the latent state\n",
    "accurately, even with fairly high noise levels, except when C=0.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Interactive Demo 4.2: LQG controller output with varying control gains\n",
    "\n",
    "Now let's implement the Kalman filter with closed-loop feedback with the controller. We will first use an arbitary control gain and a fixed value for measurement noise. We will then use the control gain from the LQR for optimal performance, with varying values for $\\rho$.\n",
    "\n",
    "(a) Visualize the system dynamics $s_t$ in closed-loop control with an arbitrary constant control gain. Vary this control gain.\n",
    "\n",
    "(b) Vary $\\rho$ to visualize the output of the optimal LQG controller. Here, we will use an optimal *constant* control gain, which is optimal in the case of an infinite time horizon (get to the goal and stay there forever)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "def simulate_kf_with_control(D=0.9, B=2., C=1., L=-0.1, T=50, ini_state=5,\n",
    "                           proc_noise = 0.1, meas_noise = 0.2):\n",
    "\n",
    "  control_gain = np.ones(T)*L\n",
    "\n",
    "  # Format the above variables into a format acccepted by the Kalman Filter\n",
    "  n_dim_state = 1\n",
    "  n_dim_obs = 1\n",
    "  n_timesteps = T\n",
    "\n",
    "  transition_matrices = np.eye(n_dim_state) * D\n",
    "  transition_covariance = np.eye(n_dim_obs) * proc_noise # process noise\n",
    "  observation_matrices = np.eye(n_dim_state) * C\n",
    "  observation_covariance = np.eye(n_dim_obs) * meas_noise\n",
    "  initial_state_mean = np.ones(n_dim_state) * ini_state\n",
    "  initial_state_covariance = np.eye(n_dim_state) * .01\n",
    "  control_matrices = np.eye(n_dim_state) * B\n",
    "\n",
    "  my_kf = MyKalmanFilter(n_dim_state, n_dim_obs, transition_matrices,\n",
    "                      transition_covariance, observation_matrices,\n",
    "                      observation_covariance, initial_state_mean,\n",
    "                      initial_state_covariance, control_matrices)\n",
    "\n",
    "  my_kf.plot_state_vs_time(n_timesteps, control_gain, goal = np.zeros(T),\n",
    "                          title='State estimation with KF (controller gain = {})'.format(L))\n",
    "\n",
    "\n",
    "widget=interactive(simulate_kf_with_control, {'manual': True},\n",
    "                   D=fixed(.9),\n",
    "                   B=fixed(2.),\n",
    "                   C=(0., 3., 1.),\n",
    "                   proc_noise=(0., 1., .1),\n",
    "                   meas_noise=(0.1, 1., .1),\n",
    "                   T=fixed(50),\n",
    "                   L=(-0.5, 0., .1),\n",
    "                   ini_state=fixed(5.))\n",
    "\n",
    "widget.children[-2].description = 'Run Simulation'\n",
    "widget.children[-2].style.button_color = 'lightgreen'\n",
    "controls = HBox(widget.children[:-1], layout=Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Interactive Demo 4.3: LQG with varying control effort costs\n",
    "\n",
    "Now let's see the performance of the LQG controller. We will use an LQG controller gain, where the control gain is from a system with an infinite-horizon. In this case, the optimal control gain turns out to be a constant. \n",
    "\n",
    "Vary the value of $\\rho$ from $0$ to large values, to see the effect on the state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Execute this cell to include helper function for LQG\n",
    "\n",
    "class LQG(MyKalmanFilter, LQR):\n",
    "  def __init__(self, T, n_dim_state, n_dim_obs,\n",
    "                transition_matrices, transition_covariance, observation_matrices,\n",
    "                observation_covariance, initial_state_mean, initial_state_covariance, control_matrices):\n",
    "    MyKalmanFilter.__init__(self,n_dim_state, n_dim_obs,\n",
    "                            transition_matrices, transition_covariance,\n",
    "                            observation_matrices,observation_covariance,\n",
    "                            initial_state_mean, initial_state_covariance, control_matrices)\n",
    "    LQR.__init__(self,T, initial_state_mean, transition_covariance)\n",
    "\n",
    "  def control_gain_LQR_infinite(self, rho):\n",
    "    control_gain_LQR_finite = self.control_gain_LQR(self.transition_matrices, self.control_matrices, rho)\n",
    "    return control_gain_LQR_finite[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "def simulate_kf_with_lqg(D=0.9, B=2., C=1., T=50, ini_state=5,\n",
    "                           proc_noise=0.1, meas_noise=0.2, rho=1.):\n",
    "\n",
    "    # Format the above variables into a format acccepted by the Kalman Filter\n",
    "    n_dim_state = 1\n",
    "    n_dim_obs = 1\n",
    "    n_timesteps = T\n",
    "\n",
    "    transition_matrices = np.eye(n_dim_state) * D\n",
    "    transition_covariance = np.eye(n_dim_obs) * proc_noise # process noise\n",
    "    observation_matrices = np.eye(n_dim_state) * C\n",
    "    observation_covariance = np.eye(n_dim_obs) * meas_noise\n",
    "    initial_state_mean = np.ones(n_dim_state) * ini_state\n",
    "    initial_state_covariance = np.eye(n_dim_state) * .01\n",
    "    control_matrices = np.eye(n_dim_state) * B\n",
    "\n",
    "    my_kf = MyKalmanFilter(n_dim_state, n_dim_obs, transition_matrices,\n",
    "                        transition_covariance, observation_matrices,\n",
    "                        observation_covariance, initial_state_mean,\n",
    "                        initial_state_covariance, control_matrices)\n",
    "\n",
    "    lqg = LQG(n_timesteps, n_dim_state, n_dim_obs,\n",
    "            transition_matrices, transition_covariance, observation_matrices,\n",
    "            observation_covariance, initial_state_mean, initial_state_covariance,\n",
    "            control_matrices)\n",
    "\n",
    "    control_gain_lqg = lqg.control_gain_LQR_infinite(rho) * np.ones(n_timesteps)\n",
    "\n",
    "    lqg.plot_state_vs_time(n_timesteps, control_gain_lqg, goal = np.zeros(T),\n",
    "                        title='State estimation with KF (LQG controller)')\n",
    "\n",
    "widget=interactive(simulate_kf_with_lqg, {'manual': True},\n",
    "                    D = fixed(.9),\n",
    "                    B = fixed(2.),\n",
    "                    C = fixed(1.),\n",
    "                    proc_noise = fixed(.1),\n",
    "                    meas_noise = fixed(.2),\n",
    "                    T = fixed(50),\n",
    "                    ini_state = fixed(5.),\n",
    "                    rho=(0., 5., 1.))\n",
    "\n",
    "widget.children[-2].description = 'Run Simulation'\n",
    "widget.children[-2].style.button_color = 'lightgreen'\n",
    "controls = HBox(widget.children[:-1], layout = Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "### Interactive Demo 4.4: How does the process noise and the measurement noise influence the controlled state and desired action?\n",
    "\n",
    "Process noise $w_t$ (proc_noise) and measurement noise $v_t$ (meas_noise) have very different effects on the controlled state. \n",
    "\n",
    "(a) To visualize this, play with the sliders to get an intuition for how process noise and measurement noise influences the controlled state. How are these two sources of noise different?\n",
    "\n",
    "(b) Next, for varying levels of process noise and measurement noise (note that the control policy is exactly the same for all these values), plot the mean squared error (MSE) between state and the goal, as well as the control cost. What do you notice?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Make sure you execute this cell to enable the widget!\n",
    "\n",
    "def lqg_slider(D=0.9, B=2., C=1., T=50, ini_state=5,\n",
    "                           proc_noise=2.9, meas_noise=0., rho=1.):\n",
    "\n",
    "    # Format the above variables into a format acccepted by the Kalman Filter\n",
    "    # Format the above variables into a format acccepted by the Kalman Filter\n",
    "    n_dim_state = 1\n",
    "    n_dim_obs = 1\n",
    "    n_timesteps = T\n",
    "\n",
    "    transition_matrices = np.eye(n_dim_state) * D\n",
    "    transition_covariance = np.eye(n_dim_obs) * proc_noise # process noise\n",
    "    observation_matrices = np.eye(n_dim_state) * C\n",
    "    observation_covariance = np.eye(n_dim_obs) * meas_noise\n",
    "    initial_state_mean = np.ones(n_dim_state) * ini_state\n",
    "    initial_state_covariance = np.eye(n_dim_state) * .01\n",
    "    control_matrices = np.eye(n_dim_state) * B\n",
    "    rho = 1\n",
    "\n",
    "    lqg = LQG(n_timesteps, n_dim_state, n_dim_obs,\n",
    "             transition_matrices, transition_covariance, observation_matrices,\n",
    "             observation_covariance, initial_state_mean, initial_state_covariance, control_matrices)\n",
    "\n",
    "    control_gain_lqg = lqg.control_gain_LQR_infinite(rho) * np.ones(n_timesteps)\n",
    "    lqg.plot_state_vs_time(n_timesteps, control_gain_lqg, goal = np.zeros(n_timesteps),\n",
    "                           title='State estimation with KF (LQG controller)')\n",
    "\n",
    "widget=interactive(lqg_slider, {'manual': True},\n",
    "                    D = fixed(.9),\n",
    "                    B = fixed(2.),\n",
    "                    C = fixed(1.),\n",
    "                    proc_noise = (0., 3., .1),\n",
    "                    meas_noise = (0.1, 3., .1),\n",
    "                    T = fixed(50),\n",
    "                    ini_state = fixed(5.),\n",
    "                    rho=fixed(1.))\n",
    "\n",
    "widget.children[-2].description = 'Run Simulation'\n",
    "widget.children[-2].style.button_color = 'lightgreen'\n",
    "controls = HBox(widget.children[:-1], layout = Layout(flex_flow='row wrap'))\n",
    "output = widget.children[-1]\n",
    "display(VBox([controls, output]));"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#to_remove explanation\n",
    "\"\"\"\n",
    "As you increase the process noise, you will notice that it becomes more\n",
    "difficult to keep the state close to the goal g=0, even though we may have very\n",
    "little measurement noise (thus can estimate the state exactly).\n",
    "\n",
    "On the other hand, as you increase the measurement noise, you will notice that\n",
    "it is harder to estimate the states, and this also may make it harder to keep the\n",
    "state close to the goal.\n",
    "\n",
    "Which has a larger effect? How does this effect the required action a[t]?\n",
    "We will quantify these in the next section.\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "## Section 4.2 Noise effects on the LQG\n",
    "\n",
    "We can now quantify how the state cost and control costs changes when we change the process and measurement noise levels. To do so, we will run many simulations, stepping through levels of process and measurement noise, tracking MSE and cost of control for each.  Run the cell below to perform this simulations and plot them. How do you interpret the results?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#@markdown Execute this cell to to quantify the dependence of state and control\n",
    "#@markdown cost on process and measurement noise (takes ~20 seconds)\n",
    "\n",
    "D = 0.9 # state parameter\n",
    "B = 2   # control parameter\n",
    "C = 1   # measurement parameter\n",
    "noise_var = 0.1\n",
    "T = 200  # time horizon\n",
    "ini_state = 5      # initial state\n",
    "process_noise_var = 0.1   # process noise\n",
    "measurement_noise_var = 0.2     # measurement noise\n",
    "rho = 1\n",
    "\n",
    "# Format the above variables into a format acccepted by the Kalman Filter\n",
    "n_dim_state = 1\n",
    "n_dim_obs = 1\n",
    "n_timesteps = T\n",
    "\n",
    "transition_matrices = np.eye(n_dim_state) * D\n",
    "transition_covariance = np.eye(n_dim_obs) * noise_var # process noise\n",
    "observation_matrices = np.eye(n_dim_state) * C\n",
    "observation_covariance = np.eye(n_dim_obs) * measurement_noise_var\n",
    "initial_state_mean = np.ones(n_dim_state) * ini_state\n",
    "initial_state_covariance = np.eye(n_dim_state) * .01\n",
    "control_matrices = np.eye(n_dim_state) * B\n",
    "\n",
    "# Implement LQG control over n_iter iterations, and record the MSE between state and goal\n",
    "MSE_array_N_meas = []\n",
    "MSE_array_N_proc = []\n",
    "Jcontrol_array_N_meas = []\n",
    "Jcontrol_array_N_proc = []\n",
    "n_iter = 10\n",
    "meas_noise_array = np.linspace(0,3,20)\n",
    "proc_noise_array = np.linspace(0.1,3,20)\n",
    "\n",
    "for i in range(n_iter):\n",
    "  MSE_array = np.zeros(proc_noise_array.shape)\n",
    "  Jcontrol_array = np.zeros(meas_noise_array.shape)\n",
    "  for i in range(len(proc_noise_array)):\n",
    "\n",
    "    transition_covariance = np.eye(n_dim_obs)  * proc_noise_array[i]\n",
    "    observation_covariance = np.eye(n_dim_obs) * measurement_noise_var\n",
    "    lqg = LQG(n_timesteps, n_dim_state, n_dim_obs,\n",
    "          transition_matrices, transition_covariance, observation_matrices,\n",
    "          observation_covariance, initial_state_mean, initial_state_covariance, control_matrices)\n",
    "\n",
    "    control_gain_lqg = lqg.control_gain_LQR_infinite(rho) * np.ones(n_timesteps)  # Get the control gain\n",
    "    filtered_state_means_impl, filtered_state_covariances_impl, latent, measurement, control = lqg.filter_control(\n",
    "        n_timesteps, control_gain_lqg)\n",
    "    MSE_array[i] = lqg.calculate_J_state(latent)\n",
    "    Jcontrol_array[i] = lqg.calculate_J_control(control)\n",
    "\n",
    "  MSE_array_N_proc.append(MSE_array)\n",
    "  Jcontrol_array_N_proc.append(Jcontrol_array)\n",
    "\n",
    "\n",
    "  MSE_array = np.zeros(meas_noise_array.shape)\n",
    "  Jcontrol_array = np.zeros(meas_noise_array.shape)\n",
    "  for i in range(len(meas_noise_array)):\n",
    "\n",
    "    observation_covariance = np.eye(n_dim_obs)  * meas_noise_array[i]\n",
    "    transition_covariance = np.eye(n_dim_obs) * noise_var\n",
    "    lqg = LQG(n_timesteps, n_dim_state, n_dim_obs,\n",
    "          transition_matrices, transition_covariance, observation_matrices,\n",
    "          observation_covariance, initial_state_mean, initial_state_covariance, control_matrices)\n",
    "\n",
    "    control_gain_lqg = lqg.control_gain_LQR_infinite(rho) * np.ones(n_timesteps)  # Get the control gain\n",
    "    filtered_state_means_impl, filtered_state_covariances_impl, latent, measurement, control = lqg.filter_control(\n",
    "        n_timesteps, control_gain_lqg)\n",
    "    MSE_array[i] = lqg.calculate_J_state(latent)\n",
    "    Jcontrol_array[i] = lqg.calculate_J_control(control)\n",
    "\n",
    "  MSE_array_N_meas.append(MSE_array)\n",
    "  Jcontrol_array_N_meas.append(Jcontrol_array)\n",
    "\n",
    "MSE_array_proc_mean = np.mean(np.array(MSE_array_N_proc), axis = 0)\n",
    "MSE_array_proc_std = np.std(np.array(MSE_array_N_proc), axis = 0)\n",
    "MSE_array_meas_mean = np.mean(np.array(MSE_array_N_meas), axis = 0)\n",
    "MSE_array_meas_std = np.std(np.array(MSE_array_N_meas), axis = 0)\n",
    "\n",
    "Jcontrol_array_proc_mean = np.mean(np.array(Jcontrol_array_N_proc), axis = 0)\n",
    "Jcontrol_array_proc_std = np.std(np.array(Jcontrol_array_N_proc), axis = 0)\n",
    "Jcontrol_array_meas_mean = np.mean(np.array(Jcontrol_array_N_meas), axis = 0)\n",
    "Jcontrol_array_meas_std = np.std(np.array(Jcontrol_array_N_meas), axis = 0)\n",
    "\n",
    "# Visualize the quantification\n",
    "f, axs = plt.subplots(2, 2, sharex=True, sharey=True, figsize=(10, 8))\n",
    "\n",
    "axs[0,0].plot(proc_noise_array, MSE_array_proc_mean, 'r-')\n",
    "axs[0,0].fill_between(proc_noise_array, MSE_array_proc_mean+MSE_array_proc_std,\n",
    "                MSE_array_proc_mean-MSE_array_proc_std, facecolor='tab:gray', alpha=0.5)\n",
    "axs[0,0].set_title('Effect of process noise')\n",
    "axs[0,0].set_ylabel('State Cost (MSE between state and goal)')\n",
    "\n",
    "axs[0,1].plot(meas_noise_array, MSE_array_meas_mean, 'r-')\n",
    "axs[0,1].fill_between(meas_noise_array, MSE_array_meas_mean+MSE_array_meas_std,\n",
    "                MSE_array_meas_mean-MSE_array_meas_std, facecolor='tab:gray', alpha=0.5)\n",
    "axs[0,1].set_title('Effect of measurement noise')\n",
    "\n",
    "axs[1,0].plot(proc_noise_array, Jcontrol_array_proc_mean, 'r-')\n",
    "axs[1,0].fill_between(proc_noise_array, Jcontrol_array_proc_mean+Jcontrol_array_proc_std,\n",
    "                Jcontrol_array_proc_mean-Jcontrol_array_proc_std, facecolor='tab:gray', alpha=0.5)\n",
    "axs[1,0].set_xlabel('Process Noise')\n",
    "axs[1,0].set_ylabel('Cost of Control')\n",
    "\n",
    "axs[1,1].plot(meas_noise_array, Jcontrol_array_meas_mean, 'r-')\n",
    "axs[1,1].fill_between(meas_noise_array, Jcontrol_array_meas_mean+Jcontrol_array_meas_std,\n",
    "                Jcontrol_array_meas_mean-Jcontrol_array_meas_std, facecolor='tab:gray', alpha=0.5)\n",
    "axs[1,1].set_xlabel('Measurement Noise')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {}
   },
   "outputs": [],
   "source": [
    "#to_remove explanation\n",
    "\"\"\"\n",
    "While both sources of noise have an effect on the controlled state, the\n",
    "process noise has a much larger effect. As the process noise w[t] increases,\n",
    "state cost (MSE between state and goal) and  control cost increase drastically.\n",
    "You can get an intuition as to why using the sliders in the demo above.  To make\n",
    "matters worse, as the process noise gets larger, you will also need to put in\n",
    "more effort to keep the system close to the goal.\n",
    "\n",
    "The measurement noise v[t]  also has an effect on the accuracy of the\n",
    "controlled state. As this noise increases, the MSE between the state and goal\n",
    "increases. The cost of control in this case remains fairly constant with\n",
    "increasing levels of measurement noise.\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "execution": {}
   },
   "source": [
    "---\n",
    "# Summary\n",
    "\n",
    "In this tutorial, you have extended the idea of optimal policy to the Astrocat example. You have learned about how to design an optimal controller with full observation of the state (linear quadratic regulator - LQR), and under partial observability of the state (linear quadratic gaussian - LQG)."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "W9_Tutorial2",
   "provenance": [],
   "toc_visible": true
  },
  "interpreter": {
   "hash": "9516f62da91337f10c2adbe814d9c63a4b08f8271333386358218606edb781e3"
  },
  "kernel": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "kernelspec": {
   "display_name": "Python 3.7.11 64-bit ('pw3': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
